import numpy as np
import cv2
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle
import logging

# 配置日志记录
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

def apply_filters(image):
    """
    应用图像滤波器
    
    Args:
        image (numpy.ndarray): 输入灰度图像
    
    Returns:
        tuple: 不同滤波器处理后的图像
    """
    # 均值滤波
    mean_filtered = cv2.blur(image, (5,5))  # 使用5x5的核进行均值滤波，减少噪声
    
    # 高斯滤波
    gaussian_filtered = cv2.GaussianBlur(image, (5,5), 0)  # 使用5x5的核进行高斯滤波，减少噪声
    
    # Sobel滤波
    sobel_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)  # 计算x方向上的Sobel梯度
    sobel_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)  # 计算y方向上的Sobel梯度
    sobel = np.sqrt(sobel_x**2 + sobel_y**2)  # 计算梯度幅值，表示边缘强度
    
    # Canny边缘检测
    canny = cv2.Canny(image, 100, 200)  # 使用Canny算法进行边缘检测，阈值为100和200
    
    logger.info("应用图像滤波器完成")
    return mean_filtered, gaussian_filtered, sobel, canny

def create_features(image):
    """
    从图像创建特征矩阵
    
    Args:
        image (numpy.ndarray): 输入灰度图像
    
    Returns:
        numpy.ndarray: 特征矩阵
    """
    # 应用滤波器
    mean_filtered, gaussian_filtered, sobel, canny = apply_filters(image)  # 调用apply_filters函数获取滤波结果
    
    # 将原图和所有滤波结果展平并组合
    features = np.column_stack((
        image.ravel(),  # 展平原始图像
        mean_filtered.ravel(),  # 展平均值滤波后的图像
        gaussian_filtered.ravel(),  # 展平高斯滤波后的图像
        sobel.ravel(),  # 展平Sobel滤波后的图像
        canny.ravel()  # 展平Canny边缘检测后的图像
    ))
    
    logger.info(f"创建的特征矩阵维度: {features.shape}")
    return features

def main():
    """
    主程序入口，协调砂岩图像分割模型的训练和保存流程
    """
    logger.info("开始砂岩图像分割模型训练和保存流程...")
    
    try:
        # 读取图像
        image = cv2.imread('Sandstone_1.tif', cv2.IMREAD_GRAYSCALE)
        segment = cv2.imread('Sandstone_1_segment.tif', cv2.IMREAD_GRAYSCALE)
        
        if image is None or segment is None:
            raise FileNotFoundError("未找到图像文件")
        
        logger.info(f"图像形状: {image.shape}")
        
        # 创建特征
        X = create_features(image)
        y = segment.ravel()
        
        logger.info("完成从砂岩截面图1及其对应分区中获取X和y")
        
        # 分割训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        logger.info("完成数据集的分割")
        
        # 训练随机森林模型
        clf = RandomForestClassifier(n_estimators=100, random_state=42)
        clf.fit(X_train, y_train)
        logger.info("完成随机森林模型clf的训练")
        
        # 计算准确率
        accuracy = clf.score(X_test, y_test)
        logger.info(f"准确率: {accuracy:.2%}")
        
        # 保存模型
        with open('sandstone_clf.pkl', 'wb') as f:
            pickle.dump(clf, f)
        logger.info("已保存随机森林模型")
        
        logger.info("砂岩图像分割模型训练和保存流程完成！")
        
    except Exception as e:
        logger.error(f"处理过程中发生错误: {str(e)}")

if __name__ == "__main__":
    main()